

import json
from PIL import Image
import numpy as np


import cv2
import argparse
from model.build_BiSeNet import BiSeNet
import os
import torch
import cv2
from imgaug import augmenters as iaa
from PIL import Image
from torchvision import transforms
import numpy as np

MODEL_PATH="/project/train/models/best_dice_loss.pth"
DEBUG=False

# 裸土
# TEST_PIC_PATH="/home/data/865/ZDSmound20220305_V1_train_292.jpg"
# CLASS_NUM=2
# 渣土
CLASS_NUM=4
TEST_PIC_PATH="demo.jpg"

def colour_code_segmentation_my(image, label_values):

    # colour_codes =np.array( [[1, 80, 1],
    #       [2, 160, 2],
    #       [3, 240, 3],
    #       [0, 0, 0]],dtype=np.uint8)
    colour_codes = np.array([[0, 0, 0],[1, 80, 1],
                             [2, 160, 2],
                             [3, 240, 3]
                             ], dtype=np.uint8)
    x = colour_codes[image.astype(int)]
    return x


class Model():
    def __init__(self):
      # basic parameters
        parser = argparse.ArgumentParser()
        parser.add_argument('--image', action='store_true',
                            default=False, help='predict on image')
        parser.add_argument('--video', action='store_true',
                            default=False, help='predict on video')
        parser.add_argument('--checkpoint_path', type=str,
                            default='/project/train/models/best_dice_loss.pth')
        parser.add_argument('--num_classes', type=int, default=2)
        parser.add_argument('--data', type=str, default=None,
                            help='Path to image or video for prediction')
        parser.add_argument('--crop_height', type=int, default=720,
                            help='Height of cropped/resized input image to network')
        parser.add_argument('--crop_width', type=int, default=960,
                            help='Width of cropped/resized input image to network')
        parser.add_argument('--cuda', type=str, default='0',
                            help='GPU ids used for training')
        parser.add_argument('--use_gpu', type=bool, default=True,
                            help='Whether to user gpu for training')
        parser.add_argument('--csv_path', type=str, default="luotu.txt")
        parser.add_argument('--save_path', type=str, default=None,
                            required=True, help='Path to save predict image')
        parser.add_argument('--context_path', type=str, default="resnet18")
        params = [
            '--image',
            '--data', 'd.jpg',
            '--cuda', '0',
            '--save_path', 'out.png',
            '--context_path', 'resnet18'
        ]
        self.args = parser.parse_args(params)

        # build model
        os.environ['CUDA_VISIBLE_DEVICES'] = "0"
        print('load model from %s ...' % MODEL_PATH)
        model = BiSeNet(CLASS_NUM, self. args.context_path)
        print("CLASS_NUM",CLASS_NUM)
        self.model = model.cuda()
        # load pretrained model if exists
        
        self.model.load_state_dict(torch.load(MODEL_PATH))
        print('Done!')
        self.model.eval()
        # predict_on_image(model, args)

    @torch.inference_mode()
    def __call__(self, image_bgr):
        # pre-processing on image
        # print("input:",args.data)
        # image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
        image_bgr = cv2.cvtColor(image_bgr, cv2.COLOR_BGR2RGB)
        src_size = image_bgr.shape
        image = cv2.resize(image_bgr, (self.args.crop_width,
                           self.args.crop_height), interpolation=cv2.INTER_CUBIC)
        image = transforms.ToTensor()(image)
        image = transforms.Normalize(
            (0.485, 0.456, 0.406), (0.229, 0.224, 0.225))(image).unsqueeze(0)

        predict = self. model(image.cuda()).squeeze()
        predict = predict.permute(1, 2, 0)  # 维度转换
        predict = torch.argmax(predict, dim=-1)
        
        predict = predict.cpu().numpy().astype(np.uint8)
        
        # predict = colour_code_segmentation_my(
        #     predict,None).astype(np.uint8)
        # predict = cv2.resize(predict, (src_size[1], src_size[0]),cv2.INTER_NEAREST)
        # if DEBUG:
        #     print(predict.shape)
        #     cv2.imshow("a", predict)
        #     cv2.waitKey()
        # print(predict.size())
        # predict =
        # # print(predict.shape)
        # predict = cv2.resize(np.uint8(predict), (src_size[1],src_size[0]))

        return predict


def init():
    model = Model()
    return model


def process_image(handle=None, input_image=None, args=None, **kwargs):
    """Do inference to analysis input_image and get output
        Attributes:
        handle: algorithm handle returned by init()
        input_image (numpy.ndarray): image to be process, format: (h, w, c), BGR
        ## 格式标准:
        args: string in JSON format, format: {
        "mask_output_path": "/path/to/output/mask.png"
        }
        Returns: process result
    """
    args = json.loads(args)
    mask_output_path = args['mask_output_path']

    dummy_data = handle(input_image)
    # b, g, r = cv2.split(dummy_data)
    # if DEBUG:
    #     cv2.imwrite(mask_output_path, r)
    # else:
    cv2.imwrite(mask_output_path, dummy_data)
    return json.dumps({'mask': mask_output_path}, indent=4)


if __name__ == '__main__':
    model = init()
    img = cv2.imread(TEST_PIC_PATH)
    args = {'mask_output_path': 'mask_out.png', }
    args = json.dumps(args)
    process_image(model, img, args)
